// Copyright (C) 2011  Andrew H. Chan
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
// 
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

#include "input_data.h"
#include "util.h"

using namespace std;

int Input_data::read_input(std::string const& file_name) {
    cout << "---Parsing data file" << endl;
    ifstream in(file_name.c_str());
    if(!in.is_open()) {
        cout << "Could not open file " << file_name << endl;
        return 1;
    }

    try {
        // read header
        in >> num_rows;
        in >> num_cols;

        // read in segment indexes
        // assume NUM_SEGS segments total
        // note, the input file indexes its columns from 1
        int startpt = 0;
        int endpt;
        for(uint i = 0; i < NUM_SEGS - 1; ++i) {
            in >> endpt; 
            --endpt;
            seg_index.push_back(pair<int, int>(startpt, endpt));
            startpt = endpt + 1;
        }
        endpt = num_cols - 1;
        seg_index.push_back(pair<int, int>(startpt, endpt));

        // consume newline
        string garbage;
        getline(in, garbage);

        // read rest of data
        // data_matrix[hap][site]
        data_matrix.resize(num_rows);
        for(uint i = 0; i < num_rows; ++i) {
            string line;
            getline(in, line);
            data_matrix[i].resize(num_cols);
            for(uint j = 0; j < num_cols; ++j) {
                if(line[j] == '0') {
                    data_matrix[i][j] = 0;
                } else {
                    if(line[j] != '1') {
                        throw;
                    }
                    data_matrix[i][j] = 1;
                }
            }
        }
    }
    catch (...) {
        cout << "Error in parsing the file." << endl;
        return 1;
    }
    in.close();
    cout << "Done parsing" << endl;

    ////////////////////////////////////////////////////////
    // preprocess data
    ////////////////////////////////////////////////////////
    // process rows
    // remove repeated rows
    cout << "---Removing duplicate rows" << endl;
    processed_data.clear();
    boost::unordered_set<cvector> row_set;
    for(uint i=0; i < data_matrix.size(); ++i) {
        if(row_set.find(data_matrix[i]) == row_set.end()) {
            row_set.insert(data_matrix[i]);
            processed_data.push_back(data_matrix[i]);
        }
    }
    // delete row_set to save memory
    row_set.clear();
    cout << "Done removing duplicate rows" << endl;

    ///////////////////////////////////////////////////
    // process columns
    // check that each column contains both 0 and 1
    // if not, then the site is not segregating, and we can ignore it
    uint num_cols_removed = 0;
    uint num_sites = processed_data[0].size();
    vector<bool> remove_col(num_sites, false);
    cout << processed_data << endl;
    for(uint ii=0; ii<processed_data.size(); ++ii) {
        cout << processed_data[ii] << endl;
    }
    for(uint i=0; i < num_sites; ++i) {
        bool contains_both = false;
        for(uint j=0; j < processed_data.size(); ++j) {
            if(processed_data[j][i] != processed_data[0][i]) {
                contains_both = true;
                break;
            }
        }
        if(!contains_both && (num_cols_removed < num_sites - 1)) {
            //cout << "removing a column that has all 0s or all 1s" << endl;
            //for(uint pp=0; pp < processed_data.size(); ++pp) {
            //    cout << processed_data[pp][i] << " ";
            //}
            //cout << endl;
            remove_col[i] = true;
            ++num_cols_removed;
        }
    }

    // remove columns with only one bit flip
    // NOTE: need to keep track of which columns are removed
    //          and add in edges later 
    for(uint i=0; i < num_sites; ++i) {
        if(remove_col[i] == false) {
            bool contains_mult_seg = false;
            uint contains_0 = 0; 
            uint contains_1 = 0;
            for(uint j=0; j < processed_data.size(); ++j) {
                if(contains_0 <= 2 && processed_data[j][i] == 0) {
                    ++contains_0;
                }
                if(contains_1 <= 2 && processed_data[j][i] == 1) {
                    ++contains_1;
                }
                if(contains_1 >= 2 && contains_0 >= 2) {
                    contains_mult_seg = true;
                    break;
                }
            }
            //if(!contains_mult_seg && (num_cols_removed < num_sites - 1)) {
            if(!contains_mult_seg) {
                //cout << "removing a column that has exactly one 1 or one 0" << endl;
                //cout << "column: " << i << endl;
                //for(uint pp=0; pp < processed_data.size(); ++pp) {
                //   cout << processed_data[pp][i] << " ";
                //}
                //cout << endl;
                remove_col[i] = true;
                ++num_cols_removed;
            }
        }
    }
    
    //remove identical columns
    //NOTE: need to keep track of how many columns are removed
    //         and weight appropriately
    for(uint i=0; i < num_sites; ++i) {
        if(!remove_col[i]) {
            for(uint j=i+1; j < num_sites; ++j) {
                bool different = false;
                if(!remove_col[j]) {
                    for(uint k=0; k < processed_data.size(); ++k) {
                        if(processed_data[k][i] != processed_data[k][j]) {
                            different = true; 
                            break;
                        }
                    }
                    if(!different) {
                        //cout << "removing an identical column" << endl;
                        //cout << "column 1: " << i << endl;
                        //for(uint pp=0; pp < processed_data.size(); ++pp) {
                        //   cout << processed_data[pp][i] << " ";
                        //}
                        //cout << endl;
                        //cout << "column 2: " << j << endl;
                        //for(uint pp=0; pp < processed_data.size(); ++pp) {
                        //   cout << processed_data[pp][j] << " ";
                        //}
                        //cout << endl;
                        remove_col[j] = true;
                        ++num_cols_removed; 
                    }
                }
            }
        }
    }

    // remove columns that are complements of others
    // this might not be applicable...
    // NOTE: need to keep track of which columns are removed
    //          and weight appropriately
    for(uint i=0; i < num_sites; ++i) {
        if(!remove_col[i]) {
            for(uint j=i+1; j < num_sites; ++j) {
                bool complement = true;
                if(!remove_col[j]){
                    for(uint k=0; k < processed_data.size(); ++k) {
                        if(processed_data[k][i] == processed_data[k][j]) {
                            complement = false; 
                            break;
                        }
                    }
                    if(complement) {
                        //cout << "removing a column that is a complement" << endl;
                        //cout << "column 1: " << i << endl;
                        //for(uint pp=0; pp < processed_data.size(); ++pp) {
                        //   cout << processed_data[pp][i] << " ";
                        //}
                        //cout << endl;
                        //cout << "column 2: " << j << endl;
                        //for(uint pp=0; pp < processed_data.size(); ++pp) {
                        //   cout << processed_data[pp][j] << " ";
                        //}
                        //cout << endl;
                        remove_col[j] = true;
                        ++num_cols_removed; 
                    }
                }
            }
        }
    }
    
    // remove marked columns
    uint new_num_sites = 0;
    for(uint i=0; i < num_sites; ++i) {
        if(!remove_col[i]) {
            ++new_num_sites;
        }
    }
    for(uint i=0; i < processed_data.size(); ++i) {
        uint p = 0;
        for(uint j=0; j < num_sites; ++j) {
            if(!remove_col[j]) {
                processed_data[i][p] = processed_data[i][j];
                ++p;
            }
        }
    }

    for(uint i=0; i < processed_data.size(); ++i) {
        processed_data[i].resize(new_num_sites);
    }

    // remove identical rows again
    row_set.clear();
    for(uint i=0; i < processed_data.size(); ++i) {
        row_set.insert(processed_data[i]);
    }
    // initialize processed_data
    processed_data.clear();
    for(boost::unordered_set<cvector>::const_iterator iter = row_set.begin(); iter != row_set.end(); ++iter) {
        processed_data.push_back(*iter);
    }
    // delete row_set to save memory
    row_set.clear();
    cout << "size of processed data after removing repeated rows: " << processed_data.size() << endl;

    // construct segmented data from processed_data
    // need to find new indexes for segments
    //int num_to_shift = 0;
    //for(uint i=0; i < NUM_SEGS; ++i) {
    //    int start = orig_seg_index[i].first - num_to_shift;
    //    for(int k=orig_seg_index[i].first; k < orig_seg_index[i].second + 1; ++k) {
    //        if(remove_col[k]) {
    //            ++num_to_shift;
    //        }
    //    }
    //    int end = orig_seg_index[i].second - num_to_shift;
    //    seg_index.push_back(pair<int,int>(start, end));
    //    if(end < start) {
    //        throw runtime_error("ERROR: some segment has no segregating sites after data preprocessing.");
    //    }
    //    cout << "seg " << i <<": " << seg_index[i] << endl;
    //}
    
    cout << "size of processed data: " << processed_data.size() << endl;
    for(uint ii=0; ii<processed_data.size(); ++ii) {
        cout << processed_data[ii] << endl;
    }

    seg_data.resize(processed_data.size());
    for(uint i=0; i<processed_data.size(); ++i) {
        seg_data[i].resize(NUM_SEGS);
        for(uint j = 0; j < NUM_SEGS; ++j) {
            uint num_sites_seg = seg_index[j].second - seg_index[j].first + 1;
            seg_data[i][j].resize(num_sites_seg);
            for(uint k = 0; k < num_sites_seg; ++k) {
                seg_data[i][j][k] = processed_data[i][seg_index[j].first + k];
            }
        }
    }

    return 0;
}
